library(tidyverse)
library(rstan)
options(mc.cores = parallel::detectCores())
rstan_options(auto_write = T)

args <- commandArgs(trailingOnly = TRUE)

if (!any(args == "stan_data")) {
  stop("Provide a data input with the keyword stan_data")
} else if (!any(args == "saveto")) {
  stop("Provide an output filename with the keyword saveto")
} else {
  stan_data_idx <- which(args == "stan_data")
  stan_data <- readRDS(args[stan_data_idx + 1])

  saveto_idx <- which(args == "saveto")
  filename <- args[saveto_idx + 1]

  if (!any(args == "seed")) {
    seed <- 1
  } else {
    seed_idx <- which(args == "seed")
    seed <- as.integer(args[seed_idx + 1])
  }


  print(Sys.time())

  system.time(m <- stan_model(file = "Scripts/Poisson_model_lasso.stan"))

  vb_fit <- vb(object = m, data = stan_data, iter = 30000, seed = seed)

  print(paste("saving to", filename))

  res <- list(vb_fit = vb_fit, data = stan_data)

  saveRDS(res, file = paste(filename, "rds", sep = "."))

  print(Sys.time())
}
